import torch
import torch.nn.functional as F
from torch.utils import data
from torchsummary import summary
from tqdm import tqdm
import matplotlib.pyplot as plt

BATCH_SIZE = 10

class MPLClassifier(torch.nn.Module):
    def __init__(self, n_input, n_output, lr=0.01):
        super(MPLClassifier, self).__init__()
        n_hidden = round(n_input * 2/3 + n_output*2/3)
        # Net structure.
        self.hidden1 = torch.nn.Linear(n_input, n_hidden)
        self.hidden2 = torch.nn.Linear(n_hidden, round(n_hidden/2))
        self.out = torch.nn.Linear(round(n_hidden/2), n_output)
        # BP functions
        self.optimizer = torch.optim.SGD(self.parameters(), lr=lr)
        self.criterion = torch.nn.CrossEntropyLoss()
        self.loss = None
        self.loss_ = None
        self.loss_curve = []

    def forward(self, x):
        x = self.hidden1(x)
        x = F.relu_(x)
        x = self.hidden2(x)
        x = F.relu_(x)
        x = self.out(x)
        return x

    def fit(self, X, y, max_iter=200):
        self.loss_curve.clear()
        X = torch.tensor(X, dtype=torch.float)
        y = torch.tensor(y, dtype=torch.long) - 1
        for epoch in range(4):
            for i in range(max_iter):
                out = self(X)
                self.loss = self.criterion(out, y)  # cross entropy
                self.loss_ = self.loss.item()
                self.loss_curve.append(self.loss_)
                self.optimizer.zero_grad()
                self.loss.backward()
                self.optimizer.step()
                plt.clf()
                plt.plot(range(len(self.loss_curve)), self.loss_curve)
                plt.pause(0.001)
                print(self.loss_)
        plt.show()

    def predict(self, X):
        X = torch.tensor(X, dtype=torch.float)
        out = self(X)
        prediction = torch.max(out, 1)[1] + 1
        return prediction.tolist()

